package task

import (
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"log"
	"net"
	"port-scanner/vars"
	"sync"
	"time"
)

type Task struct {
	IP   string
	Port int
}

func (*Task) SynScan(dstIp string, dstPort int) error {
	srcIp, srcPort, err := getLocalIPPort(dstIp)
	if err != nil {
		log.Printf("getLocalIPPort err: %v\n", err)
		return err
	}

	ips, err := net.LookupIP(dstIp)
	if err != nil {
		log.Printf("getLocalIPPort err: %v\n", err)
		return err
	}

	dstip := ips[0].To4()
	srcport := layers.TCPPort(srcPort)
	dstport := layers.TCPPort(dstPort)

	ip := &layers.IPv4{
		SrcIP:    srcIp,
		DstIP:    dstip,
		Protocol: layers.IPProtocolTCP,
	}
	tcp := &layers.TCP{
		SrcPort: srcport,
		DstPort: dstport,
		SYN:     true,
	}

	err = tcp.SetNetworkLayerForChecksum(ip)
	if err != nil {
		log.Printf("SetNetworkLayerForChecksum err: %v", err)
		return err
	}

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		ComputeChecksums: true,
		FixLengths:       true,
	}
	if err := gopacket.SerializeLayers(buf, opts, tcp); err != nil {
		log.Printf("SerializeLayers err: %v", err)
		return err
	}

	conn, err := net.ListenPacket("ip4:tcp", "0.0.0.0")
	if err != nil {
		log.Printf("net.ListenPacket err: %v", err)
		return err
	}
	defer conn.Close()

	if _, err = conn.WriteTo(buf.Bytes(), &net.IPAddr{IP: dstip}); err != nil {
		log.Printf("conn.WriteTo err: %v", err)
		return err
	}
	if err := conn.SetDeadline(time.Now().Add(3 * time.Second)); err != nil {
		log.Printf("conn.SetDeadline err: %v", err)
		return err
	}

	for true {
		b := make([]byte, 4096)
		n, addr, err := conn.ReadFrom(b)
		if err != nil {
			log.Printf("conn.ReadFrom err: %v", err)
			return err
		}

		if addr.String() == dstip.String() {
			packet := gopacket.NewPacket(b[:n], layers.LayerTypeTCP, gopacket.Default)
			if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
				tcp, _ := tcpLayer.(*layers.TCP)
				if tcp.DstPort == srcport && tcp.SYN && tcp.ACK {
					log.Printf("%v:%d is open\n", dstIp, dstport)
					return nil
				}
			}
		}
	}
	return fmt.Errorf("no")
}

func getLocalIPPort(dstip string) (net.IP, int, error) {
	serverAddr, err := net.ResolveUDPAddr("udp", dstip+":54321")
	if err != nil {
		return nil, 0, err
	}

	// 不建立真实连接，只用于确定本地IP和端口
	if con, err := net.DialUDP("udp", nil, serverAddr); err == nil {
		if udpaddr, ok := con.LocalAddr().(*net.UDPAddr); ok {
			return udpaddr.IP, udpaddr.Port, nil
		}
	}
	return nil, -1, err
}

func Run(ipList []net.IP, ports []int) {
	taskChan := make(chan Task, vars.GoroutineNum)
	var wg sync.WaitGroup

	for i := 0; i < vars.GoroutineNum; i++ {
		go func() {
			for task := range taskChan {
				err := task.SynScan(task.IP, task.Port)
				if err == nil {
					log.Printf("ip: %s port: %d is open\n", task.IP, task.Port)
					storeResult(task)
				}
				wg.Done()
			}
		}()
	}

	for _, ip := range ipList {
		for _, port := range ports {
			wg.Add(1)
			taskChan <- Task{ip.String(), port}
		}
	}

	close(taskChan)
	wg.Wait()
}

func storeResult(task Task) {
	if v, ok := vars.Result.Load(task.IP); ok {
		if ports, ok1 := v.([]int); ok1 {
			ports = append(ports, task.Port)
			vars.Result.Store(task.IP, ports)
		}
	} else {
		ports := []int{task.Port}
		vars.Result.Store(task.IP, ports)
	}
}
